import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import create_classifier

from training.models.convnext import convnext_tiny, convnext_base_in22ft1k
from training.models.efficientnet_blocks import ConvBnAct
from training.models.geb_2 import GradientEnhancedNet
from training.models.aspp_2 import DCM

__all__ = ['fasige_2']


class FASIGE(nn.Module):

    def __init__(self, num_classes, map_classes, drop_rate, pretrained):
        super(FASIGE, self).__init__()

        self.num_classes = num_classes
        self.map_classes = map_classes
        self.drop_rate = drop_rate

        # self.encoder = convnext_tiny(pretrained=pretrained)
        self.encoder = convnext_base_in22ft1k(pretrained=pretrained)
        self.num_chs = [info['num_chs'] for info in self.encoder.feature_info]  # [96, 192, 384, 768]
        assert len(self.num_chs) == 4
        self.sobel_stream = GradientEnhancedNet(self.num_chs, filter_type='sobel')
        # self.dcm = ConvBnAct(self.num_chs[3], self.num_chs[3], kernel_size=1, pad_type='same')
        self.dcm = DCM(in_channels=self.num_chs[3], atrous_rates=[1, 6, 12, 18], out_channels=self.num_chs[3])
        self.global_pool, self.classifier = create_classifier(self.num_chs[3], self.num_classes, pool_type='avg')

    def forward(self, inputs):
        features = self.encoder(inputs)
        x_sobel = self.sobel_stream(features[0], features[1], features[2])
        x_sobel += features[3]
        dcm_feat = self.dcm(x_sobel)

        x = self.global_pool(dcm_feat)
        if self.drop_rate > 0.:
            x = F.dropout(x, p=self.drop_rate, training=self.training)
        x = self.classifier(x)

        return x


def fasige_2(num_classes=2, map_classes=1, drop_rate=0., pretrained=False):
    """ Face Anti-spoofing based on Image Gradient-enhanced """
    model = FASIGE(num_classes=num_classes, map_classes=map_classes, drop_rate=drop_rate, pretrained=pretrained)
    model.default_cfg = model.encoder.default_cfg
    return model
